/* Storage - Class used to store and retrieve pieces.
   Copyright (C) 2003 Mark J. Wielaard

   This file is part of Snark.
   
   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2, or (at your option)
   any later version.
 
   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.
 
   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software Foundation,
   Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
*/

package org.klomp.snark;

import java.io.*;
import java.util.*;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

/**
 * Maintains pieces on disk. Can be used to store and retrieve pieces.
 */
public class Storage
{
  private MetaInfo metainfo;
  private long[] lengths;
  private RandomAccessFile[] rafs;
  private String[] names;

  private final StorageListener listener;

  private final BitField bitfield;
  private int needed;

  // XXX - Not always set correctly
  int piece_size;
  int pieces;

  /** The default piece size. */
  private static int MIN_PIECE_SIZE = 256*1024;
  /** The maximum number of pieces in a torrent. */
  private static long MAX_PIECES = 100*1024/20;

  /**
   * Creates a new storage based on the supplied MetaInfo.  This will
   * try to create and/or check all needed files in the MetaInfo.
   *
   * @exception IOException when creating and/or checking files fails.
   */
  public Storage(MetaInfo metainfo, StorageListener listener)
    throws IOException
  {
    this.metainfo = metainfo;
    this.listener = listener;
    needed = metainfo.getPieces();
    bitfield = new BitField(needed);
  }

  /**
   * Creates a storage from the existing file or directory together
   * with an appropriate MetaInfo file as can be announced on the
   * given announce String location.
   */
  public Storage(File baseFile, String announce, StorageListener listener)
    throws IOException
  {
    this.listener = listener;

    // Create names, rafs and lengths arrays.
    getFiles(baseFile);
    
    long total = 0;
    ArrayList lengthsList = new ArrayList();
    for (int i = 0; i < lengths.length; i++)
      {
	long length = lengths[i];
	total += length;
	lengthsList.add(new Long(length));
      }

    piece_size = MIN_PIECE_SIZE;
    pieces = (int) ((total - 1)/piece_size) + 1;
    while (pieces > MAX_PIECES)
      {
	piece_size = piece_size*2;
	pieces = (int) ((total - 1)/piece_size) +1;
      }

    // Note that piece_hashes and the bitfield will be filled after
    // the MetaInfo is created.
    byte[] piece_hashes = new byte[20*pieces];
    bitfield = new BitField(pieces);
    needed = 0;

    List files = new ArrayList();
    for (int i = 0; i < names.length; i++)
      {
	List file = new ArrayList();
	StringTokenizer st = new StringTokenizer(names[i], File.separator);
	while (st.hasMoreTokens())
	  {
	    String part = st.nextToken();
	    file.add(part);
	  }
	files.add(file);
      }

    String name = baseFile.getName();
    if (files.size() == 1)
      {
	files = null;
	lengthsList = null;
      }

    // Note that the piece_hashes are not correctly setup yet.
    metainfo = new MetaInfo(announce, baseFile.getName(), files,
			    lengthsList, piece_size, piece_hashes, total);

  }

  // Creates piece hases for a new storage.
  public void create() throws IOException
  {
    // Calculate piece_hashes
    MessageDigest digest = null;
    try
      {
	digest = MessageDigest.getInstance("SHA");
      }
    catch(NoSuchAlgorithmException nsa)
      {
	throw new InternalError(nsa.toString());
      }

    byte[] piece_hashes = metainfo.getPieceHashes();

    byte[] piece = new byte[piece_size];
    for (int i = 0; i < pieces; i++)
      {
	int length = getUncheckedPiece(i, piece, 0);
	digest.update(piece, 0, length);
	byte[] hash = digest.digest();
	for (int j = 0; j < 20; j++)
	  piece_hashes[20 * i + j] = hash[j];

	bitfield.set(i);

	if (listener != null)
	  listener.storageChecked(this, i, true);
      }

    if (listener != null)
      listener.storageAllChecked(this);

    // Reannounce to force recalculating the info_hash.
    metainfo = metainfo.reannounce(metainfo.getAnnounce());
  }

  private void getFiles(File base) throws IOException
  {
    ArrayList files = new ArrayList();
    addFiles(files, base);

    int size = files.size();
    names = new String[size];
    lengths = new long[size];
    rafs = new RandomAccessFile[size];

    int i = 0;
    Iterator it = files.iterator();
    while (it.hasNext())
      {
	File f = (File)it.next();
	names[i] = f.getPath();
	lengths[i] = f.length();
	rafs[i] = new RandomAccessFile(f, "r");
	i++;
      }
  }

  private static void addFiles(List l, File f)
  {
    if (!f.isDirectory())
      l.add(f);
    else
      {
	File[] files = f.listFiles();
	if (files == null)
	  {
	    Snark.debug("WARNING: Skipping '" + f 
			+ "' not a normal file.", Snark.WARNING);
	    return;
	  }
	for (int i = 0; i < files.length; i++)
	  addFiles(l, files[i]);
      }
  }

  /**
   * Returns the MetaInfo associated with this Storage.
   */
  public MetaInfo getMetaInfo()
  {
    return metainfo;
  }

  /**
   * How many pieces are still missing from this storage.
   */
  public int needed()
  {
    return needed;
  }

  /**
   * Whether or not this storage contains all pieces if the MetaInfo.
   */
  public boolean complete()
  {
    return needed == 0;
  }

  /**
   * The BitField that tells which pieces this storage contains.
   * Do not change this since this is the current state of the storage.
   */
  public BitField getBitField()
  {
    return bitfield;
  }

  /**
   * Creates (and/or checks) all files from the metainfo file list.
   */
  public void check() throws IOException
  {
    File base = new File(filterName(metainfo.getName()));

    List files = metainfo.getFiles();
    if (files == null)
      {
	// Create base as file.
	Snark.debug("Creating/Checking file: " + base, Snark.NOTICE);
	if (!base.createNewFile() && !base.exists())
	  throw new IOException("Could not create file " + base);

	lengths = new long[1];
	rafs = new RandomAccessFile[1];
	names = new String[1];
	lengths[0] = metainfo.getTotalLength();
	rafs[0] = new RandomAccessFile(base, "rw");
	names[0] = base.getName();
      }
    else
      {
	// Create base as dir.
	Snark.debug("Creating/Checking directory: " + base, Snark.NOTICE);
	if (!base.mkdir() && !base.isDirectory())
	  throw new IOException("Could not create directory " + base);

	List  ls = metainfo.getLengths();
	int size = files.size();
	long total = 0;
	lengths = new long[size];
	rafs = new RandomAccessFile[size];
	names = new String[size];
	for (int i = 0; i < size; i++)
	  {
	    File f = createFileFromNames(base, (List)files.get(i));
	    lengths[i] = ((Long)ls.get(i)).longValue();
	    total += lengths[i];
	    rafs[i] = new RandomAccessFile(f, "rw");
	    names[i] = f.getName();
	  }

	// Sanity check for metainfo file.
	long metalength = metainfo.getTotalLength();
	if (total != metalength)
	  throw new IOException("File lengths do not add up "
				+ total + " != " + metalength);
      }
    checkCreateFiles();
  }

  /**
   * Removes 'suspicious' characters from the give file name.
   */
  private String filterName(String name)
  {
    // XXX - Is this enough?
    return name.replace(File.separatorChar, '_');
  }

  private File createFileFromNames(File base, List names) throws IOException
  {
    File f = null;
    Iterator it = names.iterator();
    while (it.hasNext())
      {
	String name = filterName((String)it.next());
	if (it.hasNext())
	  {
	    // Another dir in the hierarchy.
	    f = new File(base, name);
	    if (!f.mkdir() && !f.isDirectory())
	      throw new IOException("Could not create directory " + f);
	    base = f;
	  }
	else
	  {
	    // The final element (file) in the hierarchy.
	    f = new File(base, name);
	    if (!f.createNewFile() && !f.exists())
	      throw new IOException("Could not create file " + f);
	  }
      }
    return f;
  }

  private void checkCreateFiles() throws IOException
  {
    // Whether we are resuming or not,
    // if any of the files already exists we assume we are resuming.
    boolean resume = false;

    // Make sure all files are available and of correct length
    for (int i = 0; i < rafs.length; i++)
      {
	long length = rafs[i].length();
	if(length == lengths[i])
	  {
	    if (listener != null)
	      listener.storageAllocated(this, length);
	    resume = true; // XXX Could dynamicly check
	  }
	else if (length == 0)
	  allocateFile(i);
	else
	  throw new IOException("File '" + names[i]
				+ "' exists, but has wrong length");
      }

    // Check which pieces match and which don't
    if (resume)
      {
	pieces = metainfo.getPieces();
	byte[] piece = new byte[metainfo.getPieceLength(0)];
	for (int i = 0; i < pieces; i++)
	  {
	    int length = getUncheckedPiece(i, piece, 0);
	    boolean correctHash = metainfo.checkPiece(i, piece, 0, length);
	    if (correctHash)
	      {
		bitfield.set(i);
		needed--;
	      }

	    if (listener != null)
	      listener.storageChecked(this, i, correctHash);
	  }
      }

    if (listener != null)
      listener.storageAllChecked(this);
  }

  private void allocateFile(int nr) throws IOException
  {
    // XXX - Is this the best way to make sure we have enough space for
    // the whole file?
    listener.storageCreateFile(this, names[nr], lengths[nr]);
    final int ZEROBLOCKSIZE = metainfo.getPieceLength(0);
    byte[] zeros = new byte[ZEROBLOCKSIZE];
    int i;
    for (i = 0; i < lengths[nr]/ZEROBLOCKSIZE; i++)
      {
	rafs[nr].write(zeros);
	if (listener != null)
	  listener.storageAllocated(this, ZEROBLOCKSIZE);
      }
    int size = (int)(lengths[nr] - i*ZEROBLOCKSIZE);
    rafs[nr].write(zeros, 0, size);
    if (listener != null)
      listener.storageAllocated(this, size);
  }


  /**
   * Closes the Storage and makes sure that all RandomAccessFiles are
   * closed. The Storage is unusable after this.
   */
  public void close() throws IOException
  {
    for (int i = 0; i < rafs.length; i++)
      {
	synchronized(rafs[i])
	  {
	    rafs[i].close();
	  }
      }
  }

  /**
   * Returns a byte array containing the requested piece or null if
   * the storage doesn't contain the piece yet.
   */
  public byte[] getPiece(int piece) throws IOException
  {
    if (!bitfield.get(piece))
      return null;

    byte[] bs = new byte[metainfo.getPieceLength(piece)];
    getUncheckedPiece(piece, bs, 0);
    return bs;
  }

  /**
   * Put the piece in the Storage if it is correct.
   *
   * @return true if the piece was correct (sha metainfo hash
   * matches), otherwise false.
   * @exception IOException when some storage related error occurs.
   */
  public boolean putPiece(int piece, byte[] bs) throws IOException
  {
    // First check if the piece is correct.
    // If we were paranoia we could copy the array first.
    int length = bs.length;
    boolean correctHash = metainfo.checkPiece(piece, bs, 0, length);
    if (listener != null)
      listener.storageChecked(this, piece, correctHash);
    if (!correctHash)
      return false;

    boolean complete;
    synchronized(bitfield)
      {
	if (bitfield.get(piece))
	  return true; // No need to store twice.
	else
	  {
	    bitfield.set(piece);
	    needed--;
	    complete = needed == 0;
	  }
      }

    long start = piece * metainfo.getPieceLength(0);
    int i = 0;
    long raflen = lengths[i];
    while (start > raflen)
      {
	i++;
	start -= raflen;
	raflen = lengths[i];
      }
    
    int written = 0;
    int off = 0;
    while (written < length)
      {
	int need = length - written;
	int len = (start + need < raflen) ? need : (int)(raflen - start);
	synchronized(rafs[i])
	  {
	    rafs[i].seek(start);
	    rafs[i].write(bs, off + written, len);
	  }
	written += len;
	if (need - len > 0)
	  {
	    i++;
	    raflen = lengths[i];
	    start = 0;
	  }
      }

    return true;
  }

  private int getUncheckedPiece(int piece, byte[] bs, int off)
    throws IOException
  {
    // XXX - copy/paste code from putPiece().
    long start = piece * metainfo.getPieceLength(0);
    int length = metainfo.getPieceLength(piece);
    int i = 0;
    long raflen = lengths[i];
    while (start > raflen)
      {
	i++;
	start -= raflen;
	raflen = lengths[i];
      }

    int read = 0;
    while (read < length)
      {
	int need = length - read;
	int len = (start + need < raflen) ? need : (int)(raflen - start);
	synchronized(rafs[i])
	  {
	    rafs[i].seek(start);
	    rafs[i].readFully(bs, off + read, len);
	  }
	read += len;
	if (need - len > 0)
	  {
	    i++;
	    raflen = lengths[i];
	    start = 0;
	  }
      }

    return length;
  }
}
